import random
import copy
import logging
import numpy as np

from util.util_log import setup_logging
from util.util_ris_pattern import point_2_phase, phase_2_pattern, phase_2_pattern_xyz, eps, \
    phase_2_bit, phaseBit_2_deg, phase_2_pattern_xyz_fft
from util.util_analysis_plane import get_peaks, get_peak_nth


# 配置日志，默认打印到控制台，也可以设置打印到文件
setup_logging()
# setup_logging(log_file="../files/multi-beam-trace/beam2/QE/L=L1+L4/trace_2025-03-19-1644-16beam-qe.txt")
# 获取日志记录器并记录日志
logger = logging.getLogger("[NSGA-II]")


# pattern 计算 peaks
def pattern_2_peaks(pattern):
    # pattern 转dB
    pattern_dbw = 20 * np.log10(np.abs(pattern) / np.max(np.max(np.abs(pattern))) + eps)
    # pattern找峰值
    peaks = get_peaks(pattern_dbw)
    return peaks


# 计算phase的方向图
def phaseBit_2_pattern(phaseBit, bit_num):
    phaseBitDeg = phaseBit_2_deg(phaseBit, bit_num)
    phaseBitRad = np.deg2rad(phaseBitDeg)
    # pattern_xyz, x, y, z = phase_2_pattern_xyz(phaseBitRad)
    pattern_xyz, x, y = phase_2_pattern_xyz_fft(phaseBitRad)  # FFT法计算, 快速
    # return np.abs(pattern_xyz), x, y
    return pattern_xyz, x, y



class Individual:
    def __init__(self, bit_num, beam_num, phaseBit, phaseBit_history_list, phaseBit_target):
        self.__beam_num = beam_num
        self.__bit_num = bit_num
        self.__phaseBit = copy.deepcopy(phaseBit)
        pattern, x, y = phaseBit_2_pattern(self.__phaseBit, self.__bit_num)
        peaks = pattern_2_peaks(pattern)
        self.__pattern = pattern
        self.__peaks = peaks
        # NSGA-II参数
        self.rank = None
        self.crowding_distance = 0
        # fitness计算参数
        self.bit_working = 1
        self.fitness_l4_map = {0: 2048, 1: 2048}
        self.phaseBitHistoryList = phaseBit_history_list
        self.phaseBitTarget = phaseBit_target
        main_lope_pos_list = self.__init_main_lope_pos(self.phaseBitTarget)
        self.main_lope_pos_list = main_lope_pos_list
        # fitness
        self.fitness = [self.fitness_l0(), self.fitness_l1(), self.fitness_l2(), self.fitness_l5()]

    def __init_main_lope_pos(self, phaseBit_target):
        # 计算目标方向图的主瓣坐标, 保存到 self.main_lope_pos_list
        pattern, x, y = phaseBit_2_pattern(phaseBit_target, self.__bit_num)
        peaks = pattern_2_peaks(pattern)
        self.main_lope_pos_list = []
        for i in range(0, self.__beam_num):
            self.main_lope_pos_list.append(peaks[i][1])
        return self.main_lope_pos_list

    # 适应度: L0 -- 主瓣检查: 检查是否是最大主瓣位置
    def __fitness_l0_pos(self, selected_peaks):
        count = 0
        for peak in selected_peaks:
            peak_x = peak[1][0]
            peak_y = peak[1][1]
            for pos in self.main_lope_pos_list:
                pos_x = pos[0]
                pos_y = pos[1]
                if abs(peak_x - pos_x) <= 1 and abs(peak_y - pos_y) <= 1:
                    count += 1
                    break
        return count

    # 适应度: L0 -- 主瓣检查: 检查main_lope_pos位置是否小于3dB
    def __fitness_l0_3dB(self, pattern):
        count = 0
        pattern_3dB = np.max(pattern) / 2.0
        for pos in self.main_lope_pos_list:
            pos_x = pos[0]
            pos_y = pos[1]
            if pattern[pos_x][pos_y] > pattern_3dB:
                count += 1
        return count

    # 适应度: L0 -- 主瓣位置检测
    def fitness_l0(self):
        # 获取前 beam_num 个 peaks
        selected_peaks = self.__peaks[:self.__beam_num]
        # count = self.__fitnes_l0_pos(selected_peaks)    # 统计这些峰值在主瓣位置列表中的数量
        count = self.__fitness_l0_3dB(self.__pattern)           # 统计主瓣位置是否小于3dB
        return np.abs(self.__beam_num - count)

    # 适应度: L1 -- PSLL 部分
    def fitness_l1(self):
        psll = 0
        # 找第beam_num+1峰(小于3dB)作为 PSLL
        if len(self.__peaks) > self.__beam_num:
            peak_nth = get_peak_nth(self.__peaks, self.__beam_num)
            if peak_nth is not None:
                psll = peak_nth[0]
        return psll

    # 适应度: L2 -- 相邻时刻码阵不同位置数量
    def fitness_l2(self):
        if self.phaseBitHistoryList is None or len(self.phaseBitHistoryList) == 0:
            return 0, 0
        phaseBitBefore = self.phaseBitHistoryList[-1]
        count = np.sum(self.__phaseBit != phaseBitBefore)
        rows, cols = self.__phaseBit.shape
        percentage = round((count / (rows * cols)) * 100, 4)
        return count, percentage

    # 适应度: L3 -- 历史码阵每个阵元工作时长均匀程度
    def fitness_l3(self):
        if self.phaseBitHistoryList is None or len(self.phaseBitHistoryList) == 0:
            return 0, 0, 0, 0, 0
        # 深拷贝 self.phaseBitHistoryList 到 phaseBit_all_list
        phaseBit_all_list = copy.deepcopy(self.phaseBitHistoryList)
        # 将 phaseBit 添加为最后一项
        phaseBit_all_list.append(self.__phaseBit)
        # 统计 phaseBit_history_list 中每个位置的值等于 target 的次数
        count_target = sum(np.array(history == self.bit_working, dtype=int) for history in phaseBit_all_list)
        # 计算所有元素的平均值
        mean_value = np.mean(count_target)
        # 计算每个元素偏离平均值的平方差，并求和
        deviation_sum = np.sum((count_target - mean_value) ** 2)
        # 计算标准差（偏差平方和开方）
        standard_deviation = np.sqrt(deviation_sum)
        # 计算归一化的标准差 (基于最大值)
        normalized_std_max = standard_deviation / np.max(count_target)
        return count_target, mean_value, deviation_sum, standard_deviation, normalized_std_max

    # 适应度: L4 -- 一段时间内码阵功率为固定值: 可以是码阵为不同值的阵元数量在一段时间内相同
    def fitness_l4(self):
        # 将二维数组展平成一维数组
        flattened = self.__phaseBit.flatten().astype(int)
        # 统计各个值的出现次数
        counts = np.bincount(flattened, minlength=len(self.fitness_l4_map))
        # 获取 fitness_l4_map 的值
        fitness_values = list(self.fitness_l4_map.values())
        # 检查 counts 和 fitness_values 的长度是否一致
        if len(counts) != len(fitness_values):
            print("Error: Length of counts and fitness_values do not match.")
            return counts, 0
        # 计算偏离对应值的方差
        variance = np.sum((counts - list(self.fitness_l4_map.values())) ** 2)
        return counts, int(variance)

    # 适应度: L5 -- 最少工作阵元, 稀疏布阵的推广
    def fitness_l5(self):
        equal_elements = np.count_nonzero(self.__phaseBit == self.bit_working)
        return equal_elements



class NSGAII:
    def __init__(self, bit_num, beam_num, population_size=50, num_generations=100, num_parents=10, mutation_rate=0.1):
        # 初始化阵列相关: 参数设置
        self.__bit_num = bit_num
        self.__beam_num = beam_num
        # 初始化遗传算法相关: 参数设置
        self.population_size = population_size
        self.num_generations = num_generations
        self.num_parents = num_parents
        self.mutation_rate = mutation_rate
        # 初始化遗传算法相关: 返回
        self.population = None
        self.best_individual = None
        self.best_fitness = None
        self.best_fitness_history = []  # 保存每一代的最佳适应度值
        self.best_individual_history = []  # 保存每一代的最佳适应度值的个体

    def initialize_population(self, phase_mix_init):
        self.population = [
            Individual(self.__bit_num, self.__beam_num, phase_mix_init, [], [])
        ] * self.population_size

    def __generate_population_part(self, before_array, current_array, max_num, phaseBit_history_list, phaseBit_target):
        part = []
        # 添加before和current作为前两个元素
        part.append(Individual(self.__bit_num, self.__beam_num, before_array, phaseBit_history_list, phaseBit_target))
        part.append(Individual(self.__bit_num, self.__beam_num, current_array, phaseBit_history_list, phaseBit_target))
        # 计算 diff 数组
        diff = before_array != current_array
        # 获取所有 False 的坐标
        false_coords = np.argwhere(diff)
        # 生成剩余的个体
        for _ in range(max_num - 2):
            # 复制 current_array
            modified_array = np.copy(current_array)
            # 随机选择一些差异行
            if np.sum(diff == True) > 0:
                # 随机选择一部分 False 坐标
                num_to_select = np.random.randint(0, len(false_coords) + 1)
                selected_false_coords = false_coords[np.random.choice(len(false_coords), num_to_select, replace=False)]
                # print("Selected False coordinates:", selected_false_coords.tolist())
                # 将 modified_array 中选取坐标替换为 before_array
                for coord in selected_false_coords:
                    row, col = coord
                    modified_array[row, col] = before_array[row, col]
                # print("Modified array:", modified_array.tolist())
            part.append(Individual(self.__bit_num, self.__beam_num, modified_array, phaseBit_history_list, phaseBit_target))
        return part

    # 初始化: ps-complex和qe一样一半, 每半第一个是ps-complex或qe, 其他是加上随机选择部分t-1和t的变化
    def initialize_population_group(self, phaseBit_mix_ps_before, phaseBit_mix_ps,
                                    phaseBit_mix_qe_before, phaseBit_mix_qe,
                                    phaseBit_history_list, phaseBit_target):
        parts_ps = self.__generate_population_part(phaseBit_mix_ps_before, phaseBit_mix_ps,
                                                   self.population_size // 2,
                                                   phaseBit_history_list, phaseBit_target)
        parts_qe = self.__generate_population_part(phaseBit_mix_qe_before, phaseBit_mix_qe,
                                                   self.population_size - self.population_size // 2,
                                                   phaseBit_history_list, phaseBit_target)
        self.population = parts_ps + parts_qe

    def non_dominated_sort(self):
        dominate_count = np.zeros(len(self.population))
        dominated_solutions = [[] for _ in range(len(self.population))]
        fronts = [[]]

        for p in range(len(self.population)):
            for q in range(len(self.population)):
                if (self.population[p].fitness[0] < self.population[q].fitness[0] and
                    self.population[p].fitness[1] < self.population[q].fitness[1]):
                    dominated_solutions[p].append(q)
                elif (self.population[q].fitness[0] < self.population[p].fitness[0] and
                      self.population[q].fitness[1] < self.population[p].fitness[1]):
                    dominate_count[p] += 1

            if dominate_count[p] == 0:
                self.population[p].rank = 0
                fronts[0].append(p)

        i = 0
        while fronts[i]:
            next_front = []
            for p in fronts[i]:
                for q in dominated_solutions[p]:
                    dominate_count[q] -= 1
                    if dominate_count[q] == 0:
                        self.population[q].rank = i + 1
                        next_front.append(q)
            i += 1
            fronts.append(next_front)

        return fronts[:-1]

    def calculate_crowding_distance(self, front):
        distance = np.zeros(len(front))
        objectives = [[], []]

        for i in front:
            objectives[0].append(self.population[i].fitness[0])
            objectives[1].append(self.population[i].fitness[1])

        for m in range(2):
            sorted_indices = np.argsort(objectives[m])
            distance[sorted_indices[0]] = distance[sorted_indices[-1]] = float('inf')
            for j in range(1, len(front) - 1):
                distance[sorted_indices[j]] += (objectives[m][sorted_indices[j + 1]] - objectives[m][sorted_indices[j - 1]])

        for i in range(len(front)):
            self.population[front[i]].crowding_distance = distance[i]

    def selection(self):
        selected = []
        fronts = self.non_dominated_sort()

        for front in fronts:
            if len(selected) + len(front) <= self.population_size:
                selected.extend(front)
            else:
                crowding_distance = [self.population[i].crowding_distance for i in front]
                sorted_front = sorted(front, key=lambda idx: self.population[idx].crowding_distance, reverse=True)
                selected.extend(sorted_front[:self.population_size - len(selected)])
                break

        return [self.population[i] for i in selected]

    def crossover(self, parent1, parent2):
        alpha = np.random.rand(*parent1.x.shape)  # 使用与 x 维度相同的 alpha
        child1 = Individual(alpha * parent1.x + (1 - alpha) * parent2.x)
        child2 = Individual((1 - alpha) * parent1.x + alpha * parent2.x)
        return child1, child2

    def mutate(self, individual):
        mutation_strength = 0.1
        individual.x += np.random.normal(0, mutation_strength, size=individual.x.shape)

    def run(self, phaseBit_mix_init, phaseBit_mix_ps_before, phaseBit_mix_ps, phaseBit_mix_qe_before, phaseBit_mix_qe,
            phaseBit_history_list, phaseBit_target):
        logger.info("population_size=%d, num_generations=%d, num_parents=%d, mutation_rate=%d"
                    % (self.population_size, self.num_generations, self.num_parents, self.mutation_rate))
        # """初始化返回值"""
        self.best_individual = Individual(self.__bit_num, self.__beam_num, phaseBit_mix_init,
                                          phaseBit_history_list, phaseBit_target)
        self.best_fitness = self.best_individual.fitness
        self.best_individual_history = []
        self.best_fitness_history = []
        # """运行遗传算法 -- 初始化阶段"""
        # self.initialize_population(phaseBit_mix_init)
        self.initialize_population_group(phaseBit_mix_ps_before, phaseBit_mix_ps,
                                         phaseBit_mix_qe_before, phaseBit_mix_qe,
                                         phaseBit_history_list, phaseBit_target)
        # """运行遗传算法 -- 搜索阶段"""
        # fixme: 从此往下的都没有改
        for generation in range(self.num_generations):
            # 选择操作
            selected_parents, selected_fitness_scores = self.selection()
            # 交换操作
            offspring = self.crossover(selected_parents, self.population_size - self.num_parents)
            # 变异操作
            offspring = self.mutation(offspring)
            # 计算后代的适应度值
            offspring_fitness_scores = [self.fitnessObj.fitness(individual) for individual in offspring]
            # 更新种群
            self.population = selected_parents + offspring
            # 合并父代和后代的适应度值
            all_fitness_scores = selected_fitness_scores + offspring_fitness_scores
            #
            # 找到当前代的最佳个体
            this_best_index = np.argmin(all_fitness_scores)  # 找到最佳适应度值的索引
            if this_best_index < len(selected_parents):
                this_best_individual = selected_parents[this_best_index]  # 最佳个体来自父代
            else:
                this_best_individual = offspring[this_best_index - len(selected_parents)]  # 最佳个体来自后代
            this_best_fitness = all_fitness_scores[this_best_index]  # 最佳适应度值
            #
            # 更新最佳适应度
            if this_best_fitness < self.best_fitness:
                self.best_fitness = this_best_fitness
                self.best_individual = this_best_individual
            # 记录最佳适应度曲线
            self.best_fitness_history.append(self.best_fitness)
            self.best_individual_history.append(self.best_individual)
            #
            logger.info("generation=%d: self.best_fitness=%f" % (generation, self.best_fitness))
            logger.info("---------------------------------------------------------------------------")
        return self.best_individual, self.best_fitness, self.best_fitness_history, self.best_individual_history


# 使用示例
if __name__ == "__main__":
    nsga2 = NSGAII(population_size=100, num_rows=3, num_cols=3, generations=100)
    final_population = nsga2.run()

    # 输出优化结果
    for individual in final_population:
        print(f"X:\n{individual.x}, f1: {individual.fitness[0]:.2f}, f2: {individual.fitness[1]:.2f}")
